from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from selenium import webdriver
import time
import os
import requests
from datetime import datetime

from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait


# 处理域名和文件操作
def write_hosts_file(v4, v6, file_path):
    # 获取所在时间并格式化为‘年月日时分’,输出位
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(f"## get host IP at {datetime.now().strftime('%Y-%m-%d %H:%M')} ##\n")
        f.write("127.0.0.1\tlocalhost\n::1\tlocalhost ip6-localhost ip6-loopback\nff02::1\tip6-allnodes\nff02::2\tip6-allrouters\n")
        f.write(v6 + '\n\n' + v4)

# 新增函数：处理IP列表并追加到目标字符串中
def process_ip_list(data_list, domain, ip_type, result_str):
    for data_temp in data_list:
        data = data_temp.text
        print(f"{domain}获取到的{ip_type}地址是\n{data}")
        list_data = data.split("\n")
        for final_data in list_data:
            if final_data == '-':
                print(f'{domain}尚未获取到有效的{ip_type}地址，跳过')
            else:
                result_str += f"\n{final_data}\t{domain}"
                print(result_str)
        print("操作成功，等待2s，操作下一个网页\n")
        time.sleep(2)
    return result_str

def main():
    if requests.get("https://dnschecker.org/").status_code != 200:
        print("https://dnschecker.org/ 不可达，5秒后自动退出程序")
        time.sleep(5)
        exit()

    # 配置浏览器
    opt = Options()
    opt.add_argument("--headless")
    opt.add_argument("--disable-gpu")
    driver_v4 = webdriver.Edge(options=opt)
    driver_v6 = webdriver.Edge(options=opt)
    print(driver_v4 == driver_v6)

    # 点击接受cookie按钮
    driver_v4.get('https://dnschecker.org/country/cn/#A/')
    button_accept_v4 = driver_v4.find_element(By.XPATH,'//*[@id="cookie_box_accept_button"]')
    button_accept_v4.click()
    driver_v6.get('https://dnschecker.org/country/cn/#AAAA/')
    button_accept_v6 = driver_v6.find_element(By.XPATH,'//*[@id="cookie_box_accept_button"]')
    button_accept_v6.click()

    # 获取输入框位置
    text_inputbox_v4 = driver_v4.find_element(By.XPATH, '//*[@id="q"]')
    text_inputbox_v6 = driver_v6.find_element(By.XPATH, '//*[@id="q"]')

    v4 = '## 这里是ipv4地址'
    v6 = '## 这里是ipv6地址'
    change = 'api.tmdb.org'
    # 获取所在路径
    path = os.getcwd()
    with open(path+"/domains.txt", "r", encoding="UTF-8") as f_domains:
        for domain in f_domains.readlines():
            domain = domain.strip()
            text_inputbox_v4.clear()
            text_inputbox_v6.clear()
            print(f"正在查询 {domain}")
            if domain == 'api.themoviedb.org':
                print(
                    f'由于输入的域名是{domain}，此域名查询的IP无效，'
                    f'所以将查询域名改为{change}，并将查询到的IP对应到{domain}')
                text_inputbox_v4.send_keys(change)
                text_inputbox_v6.send_keys(change)
            else:
                text_inputbox_v4.send_keys(domain)
                text_inputbox_v6.send_keys(domain)
            # 输入目标域名
            time.sleep(1)

            # 等待页面加载完成
            WebDriverWait(driver_v4, 10).until(EC.presence_of_element_located((By.XPATH, '//*[@id="s"]')))
            WebDriverWait(driver_v6, 10).until(EC.presence_of_element_located((By.XPATH, '//*[@id="s"]')))

            a = 1
            while a <= 3:
                try:
                    # 找到按钮并点击它
                    button_v4 = driver_v4.find_element(By.XPATH, '//*[@id="s"]')
                    button_v6 = driver_v6.find_element(By.XPATH, '//*[@id="s"]')
                    button_v4.click()
                    button_v6.click()
                    time.sleep(2)
                    a = 4
                except:
                    print(f'{domain}查询ip失败，重试{a}次')
                    a += 1
                    if a == 4:
                        print('尝试次数过多，跳过这个网页')

            # 获取数据,并判断是否存在
            data_list_v4 = driver_v4.find_elements(By.XPATH, '//*[@id="results"]/tbody/tr/td[2]')
            data_list_v6 = driver_v6.find_elements(By.XPATH, '//*[@id="results"]/tbody/tr/td[2]')

            # IPv4处理部分：
            v4 = process_ip_list(data_list_v4, domain, "v4", v4)

            # IPv6处理部分：
            v6 = process_ip_list(data_list_v6, domain, "v6", v6)

        driver_v4.quit()
        driver_v6.quit()
    f_domains.close()


    # 写入hosts文件
    write_hosts_file(v4, v6, path+'/hosts')
    print(f"写入文件完成，文件位置{path}/hosts.txt")


if __name__ == "__main__":
    main()
